# sg.py
import requests
from core.mcp import mcp

# ==============================
# 配置
# ==============================
BASE_URL = "https://agent.njqhjy.net/console/api"
LOGIN_ENDPOINT = "/login"

# 所有统计接口定义
ENDPOINTS = {
    "apps_count": "/statistics/apps",
    "models_count": "/statistics/models",
    "users_count": "/statistics/users",
    "tools_count": "/statistics/tools",
    "apps_by_tenant": "/statistics/apps-by-tenant"
}

# 登录凭证
LOGIN_DATA = {
    "email": "admin@163.com",
    "password": "Qinhuai@123"
}

# 全局 token 缓存
_auth_token = None


# ==============================
# 内部函数：获取认证 Token
# ==============================
def _get_token():
    """获取 JWT Token，自动登录（仅首次或过期时）"""
    global _auth_token
    if _auth_token:
        return _auth_token

    try:
        response = requests.post(
            f"{BASE_URL}{LOGIN_ENDPOINT}",
            json=LOGIN_DATA,
            verify=False  # ⚠️ 生产环境请配置证书验证
        )
        response.raise_for_status()
        data = response.json()
        print("登录成功，获取 access_token", data)
        _auth_token = data["data"]["access_token"]
        return _auth_token
    except Exception as e:
        raise RuntimeError(f"登录失败，无法获取 access_token: {str(e)}")


# ==============================
# 工具函数：调用各统计接口
# ==============================

def get_applications_count() -> int:
    """获取系统中应用的总数量"""
    headers = {"Authorization": f"Bearer {_get_token()}"}
    try:
        resp = requests.get(f"{BASE_URL}{ENDPOINTS['apps_count']}", headers=headers, verify=False)
        resp.raise_for_status()
        data = resp.json()
        return int(data.get("count", 0))
    except Exception as e:
        raise RuntimeError(f"获取应用总数失败: {str(e)}")


def get_models_count() -> int:
    """获取系统中AI模型的总数量"""
    headers = {"Authorization": f"Bearer {_get_token()}"}
    try:
        resp = requests.get(f"{BASE_URL}{ENDPOINTS['models_count']}", headers=headers, verify=False)
        resp.raise_for_status()
        data = resp.json()
        return int(data.get("count", 0))
    except Exception as e:
        raise RuntimeError(f"获取模型总数失败: {str(e)}")


def get_users_count() -> int:
    """获取系统中注册用户的总数量"""
    headers = {"Authorization": f"Bearer {_get_token()}"}
    try:
        resp = requests.get(f"{BASE_URL}{ENDPOINTS['users_count']}", headers=headers, verify=False)
        resp.raise_for_status()
        data = resp.json()
        return int(data.get("count", 0))
    except Exception as e:
        raise RuntimeError(f"获取用户总数失败: {str(e)}")


def get_tools_count() -> int:
    """获取系统中注册工具（Tools）的总数量"""
    headers = {"Authorization": f"Bearer {_get_token()}"}
    try:
        resp = requests.get(f"{BASE_URL}{ENDPOINTS['tools_count']}", headers=headers, verify=False)
        resp.raise_for_status()
        data = resp.json()
        return int(data.get("count", 0))
    except Exception as e:
        raise RuntimeError(f"获取工具总数失败: {str(e)}")


def get_applications_count_by_tenant() -> dict:
    """
    获取各租户（Tenant）下的应用数量分布
    返回示例: {"tenant_a": 10, "tenant_b": 5}
    """
    headers = {"Authorization": f"Bearer {_get_token()}"}
    try:
        resp = requests.get(f"{BASE_URL}{ENDPOINTS['apps_by_tenant']}", headers=headers, verify=False)
        resp.raise_for_status()
        data = resp.json()
        # 假设返回的是 { "data": { "tenant1": 3, "tenant2": 7 } } 或直接是对象
        return data.get("data", {}) if isinstance(data, dict) else {}
    except Exception as e:
        raise RuntimeError(f"获取租户应用分布失败: {str(e)}")
    



@mcp.tool(description="获取系统中所有应用的总数")
def get_app_count() -> int:
    return get_applications_count()


@mcp.tool(description="获取系统中AI模型的总数量")
def get_model_count() -> int:
    return get_models_count()


@mcp.tool(description="获取系统中注册用户的总数量")
def get_user_count() -> int:
    return get_users_count()


@mcp.tool(description="获取系统中注册工具的总数量")
def get_tool_count() -> int:
    return get_tools_count()


@mcp.tool(description="获取各租户下的应用数量分布")
def get_app_count_by_tenant() -> dict:
    return get_applications_count_by_tenant()